from GetFileLocation import *
from obspy import read,read_inventory,read_events
from distaz import *
from obspy.taup import TauPyModel

def seismomterCorrection(projectLocation):
    s = GetFileLocation(projectLocation)
    nevent = s.nevent
    ievent = 0
    for event in s.geventsLocation:
        if not os.path.exists(event + "SAC"):
            os.mkdir(event + "SAC")
        ievent = ievent + 1
        stations=s.getGenerator(s.getStationLocation,eventLocation=event)
        for station in stations:
            inv = read_inventory(station)
            msfs=s.getGenerator(s.getMseedLocation,eventLocation=event,station=station)
            for msf in msfs:
                st = read(msf)
                tr = st[0]
                station = tr.stats.station
                location = tr.stats.location
                channel = tr.stats.channel
                network = tr.stats.network
                starttime = tr.stats.starttime
                quality = tr.stats.mseed.dataquality
                print("Processing event # %d of %d: " % (ievent, nevent) + network + ' ' + station + ' ' + channel)
                tr.remove_response(inventory=inv,
                                   pre_filt=(0.01, 0.02, 2, 4),
                                   water_level=60,
                                   zero_mean=True,
                                   taper=True,
                                   taper_fraction=0.05, )
                sacName = ".".join((network, station, location, channel, quality)) + \
                          "%04d.%03d.%02d%02d%02d.SAC" % (starttime.year, starttime.julday,
                                                          starttime.hour, starttime.minute, starttime.second)
                tr.write(event + "/SAC/" + sacName, format='SAC')


def writeBasicInfo2sac(projectLoaction):
    s = GetFileLocation(projectLoaction)
    nevent =s.nevent
    ievent = 0
    for event in s.geventsLocation:
        ievent = ievent + 1
        try:
            quakeXMLPath = "/".join((event, event.split("\\")[-2] + ".xml"))
            e=read_events(quakeXMLPath,"QUAKEML").events[0]
        except FileNotFoundError:
            continue
        else:
            evlo, evla, evdp = e.origins[0].longitude, e.origins[0].latitude, e.origins[0].depth
            eventOrigin = e.origins[0].time
            stations=s.getGenerator(s.getStationLocation,eventLocation=event)
            for station in stations:
                inv = read_inventory(station)
                stla = inv[0][0].latitude
                stlo = inv[0][0].longitude
                stel = inv[0][0].elevation
                stdp = 0.0
                sacfs = s.getGenerator(s.getSacLocation,eventLocation=event,station=station)
                for sacf in sacfs:
                    st = read(sacf)
                    tr = st[0]
                    network = tr.stats.network
                    channel = tr.stats.channel
                    sta = tr.stats.station
                    tr.stats.sac.o = tr.stats.starttime - eventOrigin
                    print("Processing event # %d of %d: " % (ievent, nevent) + network + ' ' + sta + ' ' + channel)

                    tr.stats.sac.evdp = evdp / 1000
                    tr.stats.sac.evla = evla
                    tr.stats.sac.evlo = evlo
                    tr.stats.sac.stel = stel
                    tr.stats.sac.stla = stla
                    tr.stats.sac.stlo = stlo
                    tr.stats.sac.stdp = stdp
                    d = distaz(stla, stlo, evla, evlo)
                    tr.stats.sac.gcarc = d.delta
                    tr.stats.sac.az = d.az
                    tr.stats.sac.baz = d.baz
                    tr.write(sacf, format='SAC')

def writePhase(projectLocation):
    s=GetFileLocation(projectLocation)
    events=s.geventsLocation
    nevent=s.nevent
    ievent=0
    for event in events:
        stations=s.getGenerator(s.getStationLocation,eventLocation=event)
        ievent+=1
        print("Processing event# %d of %d"%(ievent,nevent))
        for station in stations:
            sacfs=s.getGenerator(s.getSacLocation,eventLocation=event,station=station)
            for sacf in sacfs:
                st = read(sacf)
                tr=st[0]
                if tr.stats.channel=="BHZ":
                    distance=tr.stats.sac.gcarc
                    evdp=tr.stats.sac.evdp
                    model = TauPyModel(model="iasp91")
                    arrivals = model.get_travel_times(source_depth_in_km=evdp, distance_in_degree=distance,
                                                                  phase_list=["P", "Pdiff"])
                    t=arrivals[0].time
                    r=arrivals[0].ray_param

                tr.stats.sac.t1 = t + tr.stats.sac.o
                tr.stats.sac.kt1 = 'P'
                tr.stats.sac.user1 = r
                tr.write(sacf, format='SAC')

def roateNE2RT(projectLocation):
    s = GetFileLocation(projectLocation)
    events = s.geventsLocation
    nevent = s.nevent
    ievent = 0
    elist=[]
    for event in events:
        if not os.path.exists(event + "SACZRT"):
            os.mkdir(event + "SACZRT")
        stations = s.getGenerator(s.getStationLocation, eventLocation=event)
        ievent += 1
        print("Processing event# %d of %d" % (ievent, nevent))
        for station in stations:
            sacfs = s.getGenerator(s.getSacLocation, eventLocation=event, station=station)
            st = read()
            st.clear()
            for sf in sacfs:
                st += read(sf)
            try:
                st.rotate("NE->RT", back_azimuth=st[0].stats.sac.baz)
            except ValueError:
                print("三分量长度不一致，无法进行旋转。")
                elist.append(station)
                continue
            else:
                for tr in st:
                    station = tr.stats.station
                    location = tr.stats.location
                    channel = tr.stats.channel
                    network = tr.stats.network
                    starttime = tr.stats.starttime

                    sacName = ".".join((network, station, location, channel)) + \
                            "%04d.%03d.%02d%02d%02d.SAC" % (starttime.year, starttime.julday,
                                                          starttime.hour, starttime.minute, starttime.second)
                    tr.write(event + "SACZRT/" + sacName, format='SAC')
    return elist








